#include "Basis.h"
#include "Cpu.h"

KSPIN_LOCK		g_CpuSpinLock;
KDPC			g_CpuDpcs[MAX_CPU];
ULONG			g_CpuBusyCount;

VOID InitializeCpu()
{
	KeInitializeSpinLock(&g_CpuSpinLock);
}

VOID DisposeCpu()
{
}

BOOLEAN BusyCpu(BusyCpuCallback busyCpuCallback, PVOID *pParams)
{
	ULONG cpuCount = KeQueryActiveProcessorCount(NULL);
	if (cpuCount == 1)
	{
		KdPrint(("executing callback in single cpu"));
		KIRQL oldIrq = KeRaiseIrqlToDpcLevel();
		busyCpuCallback(pParams);
		KeLowerIrql(oldIrq);
		return TRUE;
	}

	if (cpuCount > MAX_CPU)
	{
		KdPrint(("exceed max number of cpus"));
		return FALSE;
	}

	BOOLEAN retVal = FALSE;
	KIRQL oldIrq;
	KeAcquireSpinLock(&g_CpuSpinLock, &oldIrq);

	ULONG currentCpuIndex = KeGetCurrentProcessorNumber();
	for (ULONG i = 0; i < cpuCount; i++)
	{
		if (i == currentCpuIndex) continue;

		// make other cpus busy
		KdPrint(("making cpu %d busy", i));
		KeInitializeDpc(&g_CpuDpcs[i], FreezeCpuDpc, NULL);
		KeSetTargetProcessorDpc(&g_CpuDpcs[i], (CCHAR)i);
		KeSetImportanceDpc(&g_CpuDpcs[i], HighImportance);
		KeInsertQueueDpc(&g_CpuDpcs[i], (PVOID)i, NULL);
	}

	for (ULONG i = 0; i < 16; i++)
	{
		ULONG ulTmp = 1000000;
		while (ulTmp) ulTmp--;

		if (g_CpuBusyCount == cpuCount - 1)
		{
			KdPrint(("executing callback in cpu %d", currentCpuIndex));
			busyCpuCallback(pParams);
			retVal = TRUE;
			break;
		}
	}

	KeReleaseSpinLock(&g_CpuSpinLock, oldIrq);
	return retVal;
}

VOID FreezeCpuDpc(struct _KDPC *Dpc, PVOID DeferredContext, PVOID SystemArgument1, PVOID SystemArgument2)
{
	UNREFERENCED_PARAMETER(Dpc);
	UNREFERENCED_PARAMETER(DeferredContext);
	UNREFERENCED_PARAMETER(SystemArgument2);

	ULONG cpuIndex = (ULONG)SystemArgument1;
	KIRQL oldIrq = KeRaiseIrqlToDpcLevel();
	InterlockedIncrement((PLONG)&g_CpuBusyCount);
	KdPrint(("cpu %d is busy", cpuIndex));
	KeAcquireSpinLockAtDpcLevel(&g_CpuSpinLock);
	KeReleaseSpinLockFromDpcLevel(&g_CpuSpinLock);
	InterlockedDecrement((PLONG)&g_CpuBusyCount);
	KdPrint(("cpu %d is free", cpuIndex));
	KeLowerIrql(oldIrq);
}
